import numpy as np


class AuctionAlgorithm:
    """
    Class name: AuctionAlgorithm
    Version: Auction Algorithm for Symmetric and Asymmetric Type of Assignment Problems
    Input: a matrix Aij
    Ouput: Two forms of Pairs_current for task allocation
    Citation: Algorithm Conforming to Chapter 7 of the book "Network Optimization: Continuous and Discrete Models" by Dimitri P. Bertsekas, Massachusetts Institute of Technology.
              "A Forward Reverse Auction Algorithm for Asymmetric Assignment Problems," by D. P. Bertsekas and D. A. Castanon. Report LIDS-P-2159, Lab. for Information and Decision Systems, also Computational Optimization and Applications, Vol. 1, pp. 277-297, 1992.
    Notes: 1.The "Forward Reverse Auction Algorithm" above deals with only the condition for na >= nt, which "ImproveAuctionAlgorithm" could deal with any condition.
           2.Compared with Hungarian and All_permutation algorithm
    Usage: Create an object of class with input a matrix; then run method "EnterImproveAuctionAlgorithm" without anything
    Peng Wu, Apr 6, 2022
    Last modified: Apr 7, 2022
    """

    def __init__(self, task_matrix):
        self.task_matrix = task_matrix
        self.virtual_matrix = task_matrix
        self.Pairs_current = []    # pairs of assignment
        self.Pairs_bool = []    # boolen forms of pairs of assignment

    def EnterImproveAuctionAlgorithm(self):
        """
        1.上述拍卖算法未能解决matrix of bid列多行少，即na > nt的问题
        2.将报价矩阵置入，无论na和nt数量是否对称都可以得到分配结果
          非对称拍卖处理方式 -> 增加虚拟agents或虚拟targets的方式
        :param self.task_matrix: 原始竞拍矩阵
        :return: Two forms of Pairs_current for task allocation
        """
        Na = len(self.task_matrix)
        Nt = len(self.task_matrix[0])
        if Na <= Nt:
            self.AuctionInput(self.task_matrix)
            return self.Pairs_current, self.S_Reshape()
        elif Na > Nt:
            delta_n = Na - Nt
            virtual_tasks = []
            for i in range(delta_n):
                virtual_tasks.append(Nt + i)  # virtual targets and their ID
            virtual_bid = np.zeros([Na, delta_n])   # add virtual bid
            virtual_Aij = np.concatenate([self.virtual_matrix, virtual_bid], 1)
            virtual_Aij.tolist()
            self.AuctionInput(virtual_Aij)
            for i in range(len(self.Pairs_current) - 1, -1, -1):  # remove virtual items added
                if self.Pairs_current[i][1] in virtual_tasks:
                    self.Pairs_current.remove(self.Pairs_current[i])
            return self.Pairs_current, self.S_Reshape()

    def AuctionInput(self, Aij):
        # Part - 1 : Taking Inputs
        max_a = 0
        if len(Aij) == 0:  # For Manual Input
            n = int(input("\nNo of Persons : "))
            m = int(input("\nNo of Objects : "))
            print("\nEnter the Benefit(aij) martix Aij :")
            Aij = []
            for i in range(n):
                Aij.append([int(x) for x in input().split()])
                if max_a < max(Aij[i]):
                    max_a = max(Aij[i])
                if len(Aij[i]) != m:
                    print("\nError: Expected ", m, " values in a row")
                    return
        else:
            print("\nGiven Input -\n")
            m = len(Aij[0])
            n = len(Aij)
            if n > m:
                print("\nGiven Input is Wrong Way !\n")
                return
            for i in range(n):
                if max_a < max(Aij[i]):
                    max_a = max(Aij[i])

        # Part - 2 : Initial Assignment
        p = [0 for i in range(m)]  # price Variable for each object
        q = [max_a for i in range(n)]  # profit variable for each person

        self.Forward_Auction(Aij, p, q)
        return

    def Forward_Auction(self, Aij, p, q):
        # Part - 3 : Forward Auction
        # taking epsilon as 1/n
        e = 1 / len(q)
        unass = [i for i in range(len(q))]
        # Forward Auction Algorithm
        while unass:
            Pi = unass.pop()
            # finding the best j
            Ji = []
            for j in range(len(p)):
                Ji.append((Aij[Pi][j] - p[j]))
            j = np.argmax(Ji)
            v = Ji[j]
            Ji.remove(v)
            w = max(Ji)
            bid = p[j] + v - w + e
            # updating p's and q's
            p[j] = bid
            q[Pi] = Aij[Pi][j] - p[j]
            for i in range(len(self.Pairs_current) - 1, -1, -1):
                if (self.Pairs_current[i][1]) == j:
                    self.Pairs_current.remove(self.Pairs_current[i])
                    break
            self.Pairs_current.append([Pi, j])
            # print("\nCurrent Assignment: ", S_PrimaryAuction)

            # Selecting the person i to assign
            assigned = []  # list to  keep track of persons that are assigned
            for i in range(len(self.Pairs_current)):
                assigned.append(self.Pairs_current[i][0])
            if len(assigned) == len(q):
                if len(p) == len(q):
                    print("\nFinal Assignment for the Symmetric Problem is-")
                else:
                    self.Reverse_Auction(Aij, p, q)
                return
            unass = [person for person in [i for i in range(len(q))] if person not in assigned]

    def Reverse_Auction(self, Aij, p, q):
        # Part - 4 : Reverse Auction
        # taking epsilon as 1/n
        e = 1 / len(q)
        # finding the objects that are unassigned and with pj greater than lamda
        obj_a = []
        for i in range(len(self.Pairs_current)):
            obj_a.append(self.Pairs_current[i][1])
        # list of unassigned objects
        obj_una = [obj for obj in [i for i in range(len(p))] if obj not in obj_a]
        # unassigned objects with pj greater than lamda

        # for lamda value
        l = []
        for i in range(len(obj_a)):
            l.append(p[obj_a[i]])
        lamda = min(l)

        for i in range(len(obj_una) - 1, -1, -1):
            if p[obj_una[i]] <= lamda:
                obj_una.remove(obj_una[i])

        if len(obj_una) != 0:
            while obj_una != []:
                obj = obj_una.pop()
                # Finding the best person i
                Ij = []
                for i in range(len(q)):
                    Ij.append(Aij[i][obj] - q[i])
                i = np.argmax(Ij)
                v2 = Ij[i]
                Ij.remove(v2)
                w2 = max(Ij)

                if lamda >= v2 - e:
                    p[obj] = lamda
                else:
                    delta = min([v2 - lamda, v2 - w2 + e])
                    p[obj] = v2 - delta
                q[i] = Aij[i][obj] - p[obj]

                for j in range(len(self.Pairs_current) - 1, -1, -1):
                    if (self.Pairs_current[j][0]) == i:
                        self.Pairs_current.remove(self.Pairs_current[j])
                        break
                self.Pairs_current.append([i, obj])
        elif len(obj_una) == 0:
            print("\nFinal Assignment for the Asymmetric Problem is-")
            return

    def S_Reshape(self):
        """
        将任务分配"agent-task对"矩阵，转换成任务分配布尔矩阵
        :param self.task_matrix: 原竞拍bid矩阵
        :param self.Pairs_current: agent-task对
        :return: self.Pairs_bool 任务分配布尔矩阵
        """
        na = len(self.task_matrix)
        nt = len(self.task_matrix[0])
        self.Pairs_bool = np.zeros([na, nt])
        for i in range(0, len(self.Pairs_current)):
            self.Pairs_bool[self.Pairs_current[i][0]][self.Pairs_current[i][1]] = 1
        return self.Pairs_bool


# if __name__ == '__main__':
#
#     matri = [
#         [9999, 9999, 9999, 0, 0, 0],
#         [9999, 9999, 9999, 0, 0, 0],
#         [9999, 9999, 9999, 0, 0, 0],
#         [9999, 9999, 9999, 0, 0, 0],
#         [9999, 9999, 9999, 0, 0, 0],
#         [9999, 9999, 9999, 0, 0, 0]
#     ]
#
#     obj = AuctionAlgorithm(matri)
#     S, wins = obj.EnterImproveAuctionAlgorithm()
#     print(f"\nPairs - {S}")
#     print(f"\nS_bool - {wins}")